#include "jpegcodec.h"

#include <memory.h>
#include <cmath>

#include <iostream>

JpegCodec::JpegCodec()
{
    _quantMatrix = new uint[64];
#warning TODO: fill matrix with appropriate numbers
}

JpegCodec::~JpegCodec()
{
    delete[] _quantMatrix;
}

bool JpegCodec::encodeJpeg(const uchar *srcBuffer, uchar *destBuffer, uint lines, uint lineWidth)
{
    //1. copy to internal buffer
    _bufferLines = lines + ((lines%8) ? (8-(lines%8)) : 0);
    _bufferLineWidth = lineWidth + ((lineWidth%8) ? (8-(lineWidth%8)) : 0);
    for(uint8_t i=0; i<3; ++i)
        _planes[i] = new uchar[_bufferLines*_bufferLineWidth];

    //2. do RGB -> YCbCr transform
    uint8_t r,g,b;

    for(uint i=0; i<lines; ++i)
        for(uint j=0; j<lineWidth; ++j)
        {
            ulong idx = i;
            idx *= lineWidth;
            idx += j;

            //now, get pixel's color values R,G,B
            r = srcBuffer[idx];
            g = srcBuffer[idx+1];
            b = srcBuffer[idx+2];

            //transform them to the YCbCr color space and write to buffer
            ulong dstIdx = i;
            dstIdx *= _bufferLineWidth;
            dstIdx += j;
            _planes[0][dstIdx] = Y(r,g,b);
            _planes[1][dstIdx] = Cb(r,g,b);
            _planes[2][dstIdx] = Cr(r,g,b);
        }

    //3. extend image so that it could be divided to 8x8 blocks
    ulong idx; //index used to address internal buffers
    ulong lastLine = (_bufferLines-1)*_bufferLineWidth;
    //3.1 copy last line
    for(uint i=lines; i<_bufferLines; ++i)
    {
        for(uint j=0; j<_bufferLineWidth; ++j)
        {
            idx = i;
            idx *= _bufferLineWidth;
            idx += j;

            for(uint p=0; p<3; ++p)
                _planes[p][idx] = _planes[p][lastLine+j];
        }
    }
    //3.1 copy last column
    if(_bufferLineWidth>lineWidth)
    {
        ulong lastColumn;
        for(uint i=0; i<_bufferLines; ++i)
        {
            lastColumn = i;
            lastColumn *= _bufferLineWidth;
            lastColumn += lineWidth-1;

            idx = i;
            idx *= _bufferLineWidth;

            for(uint j=lineWidth; j<_bufferLineWidth; ++j)
            {
                idx += 1;
                for(uint p=0; p<3; ++p)
                    _planes[p][idx] = _planes[p][lastColumn];
            }
        }
    }

    //5. do DCT on 8x8 blocks for every plane & perform quantization
    ulong blockIdx;
    for(uint plane=0; plane<3; ++plane)
        for(uint blockI=0; blockI<_bufferLines/8; blockI+=8)
            for(uint blockJ=0; blockJ<_bufferLineWidth/8; blockJ+=8)
            {
                blockIdx = blockI;
                blockIdx *= _bufferLineWidth;
                blockIdx += blockJ;

                DCT(_planes[plane] + blockIdx);
                quantize(_planes[plane] + blockIdx);
            }

    //7. write data to destination buffer
    memcpy(destBuffer, _planes[0], _bufferLineWidth*_bufferLines);
    memcpy(destBuffer+_bufferLineWidth*_bufferLines,
                       _planes[1], _bufferLineWidth*_bufferLines);
    memcpy(destBuffer+_bufferLineWidth*_bufferLines*2,
                       _planes[2], _bufferLineWidth*_bufferLines);

    return true;
}

uint8_t JpegCodec::Y(uint8_t R, uint8_t G, uint8_t B)
{
    return (0.299*R) + (0.587*G) + (0.114*B);
}

uint8_t JpegCodec::Cb(uint8_t R, uint8_t G, uint8_t B)
{
    return 128 - (0.168736*R) - (0.331264*G) + (0.5*B);
}

uint8_t JpegCodec::Cr(uint8_t R, uint8_t G, uint8_t B)
{
    return 128 + (0.5*R) - (0.418688*G) - (0.081312*B);
}

float dctAlpha(uint u)
{
    if (u==0)
        return pow(8, -0.5);
    else
        return 0.5;
}

void JpegCodec::DCT(uchar *block)
{
    /* since the source pointer is uchar*,
       we need to convert it to signed block pointer
       to manipulate numbers in [128..127] */
    int8_t* blockSigned = reinterpret_cast<int8_t*>(block);

    double g[16];
    for(uint i=0; i<16; ++i)
        g[i] = double(blockSigned[i]-128)/128.0;

    double sum;
    for(uint u=0; u<8; ++u)
        for(uint v=0; v<8; ++v)
        {
            sum=0;
            for(uint x=0; x<8; ++x)
                for(uint y=0; y<8; ++y)
                    sum += dctAlpha(u)*dctAlpha(v) * g[y*16+x] * cos((x+0.5)*u*M_PI/8.0) * cos((y+0.5)*v*M_PI/8.0);

            blockSigned[u+(v*16)] = int8_t(sum);
        }
}

void JpegCodec::quantize(uchar *block)
{
#warning TODO
}
